package com.example.demo.aspect;

import java.io.IOException;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;

import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import com.alibaba.fastjson.JSON;
import com.example.demo.model.Result;
import com.google.common.util.concurrent.RateLimiter;

/**
 * @author jiangwz
 * @create 2019/10/27.
 */

@Component
@Aspect
public class RateLimitAspect {

    @Autowired
    private HttpServletResponse response;

    private RateLimiter rateLimiter = RateLimiter.create(1.0);//1s访问一次

    @Pointcut(value = "@annotation(com.example.demo.annotation.RateLimit)")
    public void limit(){

    }

    @Around("limit()")
    public Object around(ProceedingJoinPoint joinPoint){
        boolean acquire = rateLimiter.tryAcquire();
        Object result = null;
        try {
            if (acquire) {
                result = joinPoint.proceed();
            }else {
                output(response, JSON.toJSONString(Result.error("访问速率过快,请稍后重试")));
            }
        } catch (Throwable e) {
            e.printStackTrace();
        }
        return result;

    }

    private void output(HttpServletResponse response, String msg) throws IOException {
        response.setContentType("application/json;charset=UTF-8");
        ServletOutputStream outputStream = null;
        try {
            outputStream = response.getOutputStream();
            outputStream.write(msg.getBytes("UTF-8"));
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            outputStream.flush();
            outputStream.close();
        }
    }
}
